# Windows PowerShell Script: Batch export and edit Conda environments
# Create output directory
$outputDir = "conda_environments"
if (-not (Test-Path $outputDir)) {
    New-Item -ItemType Directory -Path $outputDir | Out-Null
}

# Get all environment list
$envs = conda env list | Where-Object { $_ -notmatch "^#" -and $_.Trim() -ne "" } | ForEach-Object {
    $parts = $_ -split '\s+', 2
    if ($parts.Count -ge 1 -and $parts[0] -ne "") {
        $parts[0]
    }
}

Write-Host "Found the following Conda environments:"
$envs | ForEach-Object { Write-Host "- $_" }
Write-Host ""

# Export each environment
foreach ($env in $envs) {
    Write-Host "Exporting environment: $env"
    
    # Export full environment
    $fullOutputPath = Join-Path $outputDir "${env}_full.yml"
    conda env export --name $env > $fullOutputPath
    
    # Export cross-platform compatible version
    $compatOutputPath = Join-Path $outputDir "${env}_compatible.yml"
    conda env export --name $env --no-builds > $compatOutputPath
    
    # Export requirements.txt (包括conda和pip包)
    $requirementsPath = Join-Path $outputDir "${env}_requirements.txt"
    Write-Host "Exporting packages to requirements.txt..."
    
    # 创建临时文件
    $tempFile = [System.IO.Path]::GetTempFileName()
    
    # 添加额外的索引URL到临时文件顶部
    @"
--extra-index-url https://mirrors.aliyun.com/pypi/simple/
--extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple/
"@ | Out-File -FilePath $tempFile -Encoding utf8
    
    # 获取conda包列表（包括python版本）
    $condaPackages = conda list --name $env | Where-Object { $_ -notmatch "^#" -and $_.Trim() -ne "" } | ForEach-Object {
        $parts = $_ -split '\s+'
        if ($parts.Count -ge 3) {
            $packageName = $parts[0]
            $version = $parts[1]
            
            # 排除一些不需要的系统包
            $excludePackages = @("_go_select", "_openmp_mutex", "libgcc", "libgomp", "libstdcxx", "libwinpthread", "ucrt", "vc", "vc14_runtime", "vs2015_runtime")
            if ($excludePackages -notcontains $packageName) {
                "$packageName==$version"
            }
        }
    }
    
    # 将conda包添加到临时文件
    $condaPackages | Out-File -FilePath $tempFile -Encoding utf8 -Append
    
    # 使用pip freeze获取pip包
    $pipCommand = "cmd.exe /c `"conda activate $env && pip freeze >> `"$tempFile`"`""
    try {
        Invoke-Expression $pipCommand
        
        if (Test-Path $tempFile) {
            # 读取临时文件内容
            $requirements = Get-Content $tempFile
            $processedRequirements = @()
            
            # 添加额外的索引URL
            $extraIndexUrls = @(
                "--extra-index-url https://mirrors.aliyun.com/pypi/simple/"
                "--extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple/"
            )
            
            # 将额外索引URL添加到文件顶部
            $processedRequirements += $extraIndexUrls
            
            foreach ($line in $requirements) {
                # 跳过空行、注释和已添加的索引URL
                if ([string]::IsNullOrWhiteSpace($line) -or $line.StartsWith("#") -or $line.StartsWith("--extra-index-url")) {
                    continue
                }
                
                # 处理已经符合格式的行
                if ($line -match "^([a-zA-Z0-9_\-\.]+)(==|>=)(.+)$") {
                    $processedRequirements += $line
                    continue
                }
                
                # 处理需要转换的行
                if ($line -match "^([a-zA-Z0-9_\-\.]+)=(.+)$") {
                    $packageName = $matches[1]
                    $version = $matches[2]
                    $processedRequirements += "$packageName==$version"
                }
                elseif ($line -match "^([a-zA-Z0-9_\-\.]+)$") {
                    $packageName = $matches[1]
                    # 对于没有版本的包，使用>=0.0.0格式
                    $processedRequirements += "$packageName>=0.0.0"
                }
                else {
                    # 保留其他格式的行
                    $processedRequirements += $line
                }
            }
            
            # 去除重复的包（可能conda和pip都有）
            $uniquePackages = @{}
            $finalRequirements = @()
            
            foreach ($line in $processedRequirements) {
                if ($line.StartsWith("--")) {
                    $finalRequirements += $line
                    continue
                }
                
                if ($line -match "^([a-zA-Z0-9_\-\.]+)(==|>=)") {
                    $packageName = $matches[1].ToLower()
                    if (-not $uniquePackages.ContainsKey($packageName)) {
                        $uniquePackages[$packageName] = $true
                        $finalRequirements += $line
                    }
                } else {
                    $finalRequirements += $line
                }
            }
            
            # 将处理后的内容写回文件
            Set-Content -Path $requirementsPath -Value $finalRequirements
            
            Write-Host "Requirements.txt file has been formatted to use package==version or package>=version format."
            
            # 检查是否有内容
            $contentCount = (Get-Content $requirementsPath | Where-Object { -not [string]::IsNullOrWhiteSpace($_) -and -not $_.StartsWith("--") }).Count
            if ($contentCount -eq 0) {
                Write-Host "Warning: No packages found for environment $env. Adding at least Python version." -ForegroundColor Yellow
                
                # 获取Python版本
                $pythonVersion = conda list --name $env python | Where-Object { $_ -notmatch "^#" -and $_.Trim() -ne "" } | ForEach-Object {
                    $parts = $_ -split '\s+'
                    if ($parts.Count -ge 3 -and $parts[0] -eq "python") {
                        $version = $parts[1]
                        "python==$version"
                    }
                }
                
                if ($pythonVersion) {
                    Add-Content -Path $requirementsPath -Value $pythonVersion
                }
            }
        } else {
            Write-Host "Warning: Failed to create requirements.txt file for $env" -ForegroundColor Yellow
        }
    } catch {
        Write-Host "Error exporting packages for $env. Please check manually." -ForegroundColor Red
    }
    
    # 清理临时文件
    if (Test-Path $tempFile) {
        Remove-Item $tempFile -Force
    }
    
    Write-Host "Exported to: $fullOutputPath, $compatOutputPath, and $requirementsPath"
    Write-Host "--------------------------------------------------"
}

Write-Host "`nAll environments have been exported to $outputDir directory"
Write-Host "You can use the following commands to rebuild environments:"
Write-Host "1. From YAML file: conda env create -f [environment_file_path].yml"
Write-Host "2. From requirements.txt:"
Write-Host "   conda create -n [env_name] python=[version]"
Write-Host "   conda activate [env_name]"
Write-Host "   pip install -r [requirements_file_path].txt"